#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File     : optimizer.py
@Project  : pipecoco
@Date     : 2021/8/15
@Author   : Zhang Jinyang
@Contact  : zhang-jy@sjtu.edu.cn
'''

from comm.config import *

class optimizer:

    def __init__(self):

        self.last_block_transmit_time = {}


    def find_optimal_config(self, model_name, record):
        last_block_transmit_time = record[3][-1] - record[2][-1]
        total_time = self.cal_pipecoco_time(record, last_block_transmit_time)
        best_performance = total_time
        print(2, 1, total_time)

        for F in range(F_scope[0], min(self.F_max[model_name], F_scope[1])):
            for b in range(2, 5):
                record = self.record[model_name]['{}_{}'.format(F, b)]
                last_block_transmit_time = self.last_block_transmit_time[model_name]['{}_{}'.format(F, b)]
                total_time = self.cal_pipecoco_time(record, last_block_transmit_time)
                print(F, b, total_time, record[-1])
                if total_time < best_performance:
                    best_performance = total_time
                    fused_layers_scope = [0, F]
                    block_num = b
        print(best_performance, fused_layers_scope, block_num)
        return fused_layers_scope, block_num


    def cal_pipecoco_time(self,record,last_block_transmit_time):

        """
        传算并行时间可根据以下公式计算出
        客户端结束时间 + 最后一块数据传输时间 + 服务端剩余计算时间
        """
        client_cal_time = record[2][-1] - record[0][0]
        server_cal_time = record[5][-1] - record[4][-1] + record[7] - record[6]
        total_time = client_cal_time + last_block_transmit_time + server_cal_time
        return total_time